{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(graph-example)=\n",
    "# Advanced model serving graph - notebook example\n",
    "\n",
    "This example demonstrates how to use MLRun serving graphs and their advanced functionality including:\n",
    "* Use of flow, task, model, and ensemble router states\n",
    "* Build tasks from custom handlers, classes and storey components\n",
    "* Use custom error handlers\n",
    "* Test graphs locally\n",
    "* Deploy the graph as a real-time serverless functions\n",
    "\n",
    "**In this example**\n",
    "- [Define functions and classes used in the graph](#define-functions-and-classes-used-in-the-graph)\n",
    "- [Create a new serving function and graph](#create-a-new-serving-function-and-graph)\n",
    "- [Test the function locally](#test-the-function-locally)\n",
    "- [Deploy the graph as a real-time serverless function](#deploy-the-graph-as-a-real-time-serverless-function)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define functions and classes used in the graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from cloudpickle import load\n",
    "from typing import List\n",
    "from sklearn.datasets import load_iris\n",
    "import numpy as np\n",
    "\n",
    "# model serving class example\n",
    "class ClassifierModel(mlrun.serving.V2ModelServer):\n",
    "    def load(self):\n",
    "        \"\"\"load and initialize the model and/or other elements\"\"\"\n",
    "        model_file, extra_data = self.get_model(\".pkl\")\n",
    "        self.model = load(open(model_file, \"rb\"))\n",
    "\n",
    "    def predict(self, body: dict) -> List:\n",
    "        \"\"\"Generate model predictions from sample.\"\"\"\n",
    "        feats = np.asarray(body[\"inputs\"])\n",
    "        result: np.ndarray = self.model.predict(feats)\n",
    "        return result.tolist()\n",
    "\n",
    "\n",
    "# echo class, custom class example\n",
    "class Echo:\n",
    "    def __init__(self, context, name=None, **kw):\n",
    "        self.context = context\n",
    "        self.name = name\n",
    "        self.kw = kw\n",
    "\n",
    "    def do(self, x):\n",
    "        print(\"Echo:\", self.name, x)\n",
    "        return x\n",
    "\n",
    "\n",
    "# error echo function, demo catching error and using custom function\n",
    "def error_catcher(x):\n",
    "    x.body = {\"body\": x.body, \"origin_state\": x.origin_state, \"error\": x.error}\n",
    "    print(\"EchoError:\", x)\n",
    "    return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# mark the end of the code section, DO NOT REMOVE !\n",
    "# mlrun: end-code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a new serving function and graph\n",
    "Use `code_to_function` to convert the above code into a serving function object and initialize a graph with async flow topology."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "function = mlrun.code_to_function(\n",
    "    \"advanced\", kind=\"serving\", image=\"mlrun/mlrun\", requirements=[\"storey\"]\n",
    ")\n",
    "graph = function.set_topology(\"flow\", engine=\"async\")\n",
    "# function.verbose = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Specify the sklearn models that are used in the ensemble."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "models_path = \"https://s3.wasabisys.com/iguazio/models/iris/model.pkl\"\n",
    "path1 = models_path\n",
    "path2 = models_path"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Build and connect the graph (DAG) using the custom function and classes and plot the result. Add states \n",
    "using the `state.to()` method (adds a new state after the current one), or using the \n",
    "`graph.add_step()` method.\n",
    "\n",
    "Use the graph `error_handler` if you want an error from the graph or a step to be fed into a specific state (catcher). See the full description in {ref}`pipelines-error-handling`.\n",
    "\n",
    "You can specify which state is the responder (returns the HTTP response) using the `state.respond()` method.\n",
    "If you don't specify the responder, the graph is non-blocking."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<!-- Generated by graphviz version 2.40.1 (20161225.0304)\n",
       " -->\n",
       "<!-- Title: mlrun&#45;flow Pages: 1 -->\n",
       "<svg width=\"636pt\" height=\"204pt\"\n",
       " viewBox=\"0.00 0.00 636.27 204.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 200)\">\n",
       "<title>mlrun&#45;flow</title>\n",
       "<polygon fill=\"#ffffff\" stroke=\"transparent\" points=\"-4,4 -4,-200 632.2712,-200 632.2712,4 -4,4\"/>\n",
       "<g id=\"clust1\" class=\"cluster\">\n",
       "<title>cluster_ensemble</title>\n",
       "<polygon fill=\"none\" stroke=\"#000000\" points=\"382.4803,-46 382.4803,-152 620.2712,-152 620.2712,-46 382.4803,-46\"/>\n",
       "</g>\n",
       "<!-- _start -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>_start</title>\n",
       "<polygon fill=\"#d3d3d3\" stroke=\"#000000\" points=\"38.5476,-120.0493 40.698,-120.1479 42.8263,-120.2953 44.9236,-120.4913 46.9815,-120.7353 48.9917,-121.0266 50.9463,-121.3645 52.8377,-121.7479 54.6587,-122.1759 56.4025,-122.6472 58.0628,-123.1606 59.634,-123.7147 61.1107,-124.308 62.4882,-124.9388 63.7625,-125.6054 64.9302,-126.3059 65.9882,-127.0385 66.9343,-127.8012 67.7669,-128.5918 68.4849,-129.4082 69.0878,-130.2481 69.5758,-131.1093 69.9496,-131.9894 70.2102,-132.886 70.3595,-133.7965 70.3997,-134.7186 70.3334,-135.6497 70.1636,-136.5873 69.8937,-137.5287 69.5276,-138.4713 69.0691,-139.4127 68.5225,-140.3503 67.8923,-141.2814 67.1831,-142.2035 66.3996,-143.114 65.5464,-144.0106 64.6285,-144.8907 63.6504,-145.7519 62.617,-146.5918 61.5329,-147.4082 60.4024,-148.1988 59.2299,-148.9615 58.0197,-149.6941 56.7755,-150.3946 55.5012,-151.0612 54.2002,-151.692 52.8757,-152.2853 51.5309,-152.8394 50.1684,-153.3528 48.7908,-153.8241 47.4003,-154.2521 45.9989,-154.6355 44.5886,-154.9734 43.1708,-155.2647 41.7472,-155.5087 40.3189,-155.7047 38.8872,-155.8521 37.4531,-155.9507 36.0175,-156 34.5815,-156 33.146,-155.9507 31.7119,-155.8521 30.2801,-155.7047 28.8519,-155.5087 27.4282,-155.2647 26.0105,-154.9734 24.6001,-154.6355 23.1988,-154.2521 21.8083,-153.8241 20.4306,-153.3528 19.0681,-152.8394 17.7233,-152.2853 16.3989,-151.692 15.0979,-151.0612 13.8236,-150.3946 12.5794,-149.6941 11.3691,-148.9615 10.1967,-148.1988 9.0662,-147.4082 7.982,-146.5918 6.9486,-145.7519 5.9706,-144.8907 5.0526,-144.0106 4.1995,-143.114 3.4159,-142.2035 2.7067,-141.2814 2.0765,-140.3503 1.53,-139.4127 1.0715,-138.4713 .7053,-137.5287 .4355,-136.5873 .2657,-135.6497 .1993,-134.7186 .2395,-133.7965 .3888,-132.886 .6495,-131.9894 1.0232,-131.1093 1.5112,-130.2481 2.1141,-129.4082 2.8321,-128.5918 3.6647,-127.8012 4.6109,-127.0385 5.6689,-126.3059 6.8365,-125.6054 8.1108,-124.9388 9.4884,-124.308 10.9651,-123.7147 12.5362,-123.1606 14.1966,-122.6472 15.9404,-122.1759 17.7614,-121.7479 19.6528,-121.3645 21.6074,-121.0266 23.6176,-120.7353 25.6755,-120.4913 27.7728,-120.2953 29.901,-120.1479 32.0515,-120.0493 34.2154,-120 36.3837,-120 38.5476,-120.0493\"/>\n",
       "<text text-anchor=\"middle\" x=\"35.2995\" y=\"-134.3\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">start</text>\n",
       "</g>\n",
       "<!-- enrich -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>enrich</title>\n",
       "<ellipse fill=\"none\" stroke=\"#000000\" cx=\"146.8955\" cy=\"-138\" rx=\"40.0939\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"146.8955\" y=\"-134.3\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">enrich</text>\n",
       "</g>\n",
       "<!-- _start&#45;&gt;enrich -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>_start&#45;&gt;enrich</title>\n",
       "<path fill=\"none\" stroke=\"#000000\" d=\"M69.9975,-138C78.2752,-138 87.2796,-138 96.1148,-138\"/>\n",
       "<polygon fill=\"#000000\" stroke=\"#000000\" points=\"96.3553,-141.5001 106.3553,-138 96.3552,-134.5001 96.3553,-141.5001\"/>\n",
       "</g>\n",
       "<!-- pre&#45;process -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>pre&#45;process</title>\n",
       "<ellipse fill=\"none\" stroke=\"#000000\" cx=\"288.8361\" cy=\"-138\" rx=\"65.7887\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"288.8361\" y=\"-134.3\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">pre&#45;process</text>\n",
       "</g>\n",
       "<!-- enrich&#45;&gt;pre&#45;process -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>enrich&#45;&gt;pre&#45;process</title>\n",
       "<path fill=\"none\" stroke=\"#000000\" d=\"M187.57,-138C195.5188,-138 204.0918,-138 212.767,-138\"/>\n",
       "<polygon fill=\"#000000\" stroke=\"#000000\" points=\"212.9665,-141.5001 222.9665,-138 212.9664,-134.5001 212.9665,-141.5001\"/>\n",
       "</g>\n",
       "<!-- catcher -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>catcher</title>\n",
       "<ellipse fill=\"none\" stroke=\"#000000\" cx=\"454.2771\" cy=\"-178\" rx=\"46.2923\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"454.2771\" y=\"-174.3\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">catcher</text>\n",
       "</g>\n",
       "<!-- pre&#45;process&#45;&gt;catcher -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>pre&#45;process&#45;&gt;catcher</title>\n",
       "<path fill=\"none\" stroke=\"#000000\" stroke-dasharray=\"5,2\" d=\"M338.4738,-150.0013C359.495,-155.0838 383.9289,-160.9913 404.9196,-166.0664\"/>\n",
       "<polygon fill=\"#000000\" stroke=\"#000000\" points=\"404.375,-169.5355 414.9175,-168.4837 406.0201,-162.7316 404.375,-169.5355\"/>\n",
       "</g>\n",
       "<!-- ensemble -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>ensemble</title>\n",
       "<polygon fill=\"none\" stroke=\"#000000\" points=\"513.8842,-90.5442 513.8842,-105.4558 478.9672,-116 429.5871,-116 394.67,-105.4558 394.67,-90.5442 429.5871,-80 478.9672,-80 513.8842,-90.5442\"/>\n",
       "<polygon fill=\"none\" stroke=\"#000000\" points=\"517.8715,-87.5709 517.8715,-108.4291 479.5543,-120 428.9999,-120 390.6827,-108.4291 390.6827,-87.5709 428.9999,-76 479.5543,-76 517.8715,-87.5709\"/>\n",
       "<text text-anchor=\"middle\" x=\"454.2771\" y=\"-94.3\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">ensemble</text>\n",
       "</g>\n",
       "<!-- pre&#45;process&#45;&gt;ensemble -->\n",
       "<g id=\"edge6\" class=\"edge\">\n",
       "<title>pre&#45;process&#45;&gt;ensemble</title>\n",
       "<path fill=\"none\" stroke=\"#000000\" d=\"M338.4738,-125.9987C354.6253,-122.0936 372.7914,-117.7015 389.8624,-113.5741\"/>\n",
       "<polygon fill=\"#000000\" stroke=\"#000000\" points=\"390.8661,-116.9323 399.7635,-111.1802 389.221,-110.1283 390.8661,-116.9323\"/>\n",
       "</g>\n",
       "<!-- ensemble/m1 -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>ensemble/m1</title>\n",
       "<ellipse fill=\"none\" stroke=\"#000000\" cx=\"585.2712\" cy=\"-72\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"585.2712\" y=\"-68.3\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">m1</text>\n",
       "</g>\n",
       "<!-- ensemble&#45;&gt;ensemble/m1 -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>ensemble&#45;&gt;ensemble/m1</title>\n",
       "<path fill=\"none\" stroke=\"#000000\" d=\"M513.6418,-86.2172C525.7485,-83.8142 538.1371,-81.3553 549.1335,-79.1727\"/>\n",
       "<polygon fill=\"#000000\" stroke=\"#000000\" points=\"550.0679,-82.5556 559.1952,-77.1756 548.7051,-75.6895 550.0679,-82.5556\"/>\n",
       "</g>\n",
       "<!-- ensemble/m2 -->\n",
       "<g id=\"node7\" class=\"node\">\n",
       "<title>ensemble/m2</title>\n",
       "<ellipse fill=\"none\" stroke=\"#000000\" cx=\"585.2712\" cy=\"-126\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"585.2712\" y=\"-122.3\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">m2</text>\n",
       "</g>\n",
       "<!-- ensemble&#45;&gt;ensemble/m2 -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>ensemble&#45;&gt;ensemble/m2</title>\n",
       "<path fill=\"none\" stroke=\"#000000\" d=\"M512.1143,-110.3627C524.8132,-113.0771 537.9014,-115.8747 549.4343,-118.3399\"/>\n",
       "<polygon fill=\"#000000\" stroke=\"#000000\" points=\"549.0219,-121.8307 559.5326,-120.4984 550.4852,-114.9853 549.0219,-121.8307\"/>\n",
       "</g>\n",
       "<!-- final -->\n",
       "<g id=\"node8\" class=\"node\">\n",
       "<title>final</title>\n",
       "<ellipse fill=\"none\" stroke=\"#000000\" cx=\"585.2712\" cy=\"-18\" rx=\"31.3957\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"585.2712\" y=\"-14.3\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">final</text>\n",
       "</g>\n",
       "<!-- ensemble&#45;&gt;final -->\n",
       "<g id=\"edge7\" class=\"edge\">\n",
       "<title>ensemble&#45;&gt;final</title>\n",
       "<path fill=\"none\" stroke=\"#000000\" d=\"M539.419,-46.0026C544.8877,-42.6628 550.155,-39.446 555.048,-36.4577\"/>\n",
       "<polygon fill=\"#000000\" stroke=\"#000000\" points=\"556.9094,-39.4221 563.6195,-31.223 553.261,-33.4481 556.9094,-39.4221\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<graphviz.dot.Digraph at 0x7fe03f6941d0>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# use built-in storey class or our custom Echo class to create and link Task steps. Add an error handling step that runs if the grah step fails\n",
    "graph.to(\"storey.Extend\", name=\"enrich\", _fn='({\"tag\": \"something\"})').to(\n",
    "    class_name=\"Echo\", name=\"pre-process\", some_arg=\"abc\"\n",
    ").error_handler(name=\"catcher\", handler=\"handle_error\", full_event=True)\n",
    "\n",
    "# add an Ensemble router with two child models (routes). The \"*\" prefix mark it is a router class\n",
    "router = graph.add_step(\n",
    "    \"*mlrun.serving.VotingEnsemble\", name=\"ensemble\", after=\"pre-process\"\n",
    ")\n",
    "router.add_route(\"m1\", class_name=\"ClassifierModel\", model_path=path1)\n",
    "router.add_route(\"m2\", class_name=\"ClassifierModel\", model_path=path2)\n",
    "\n",
    "# add the final step (after the router) that handles post processing and responds to the client\n",
    "graph.add_step(class_name=\"Echo\", name=\"final\", after=\"ensemble\").respond()\n",
    "\n",
    "# plot the graph (using Graphviz) and run a test\n",
    "graph.plot(rankdir=\"LR\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test the function locally\n",
    "Create a test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "\n",
    "iris = load_iris()\n",
    "x = random.sample(iris[\"data\"].tolist(), 5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a mock server (simulator) and test the graph with the test data.\n",
    "\n",
    "> Note: The model and router objects support a common serving protocol API, see the [protocol and API section](model-api.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> 2021-01-09 22:49:26,365 [info] model m1 was loaded\n",
      "> 2021-01-09 22:49:26,493 [info] model m2 was loaded\n",
      "> 2021-01-09 22:49:26,494 [info] Loaded ['m1', 'm2']\n",
      "Echo: pre-process {'inputs': [[6.9, 3.2, 5.7, 2.3], [6.4, 2.7, 5.3, 1.9], [4.9, 3.1, 1.5, 0.1], [7.3, 2.9, 6.3, 1.8], [5.4, 3.7, 1.5, 0.2]], 'tag': 'something'}\n",
      "Echo: final {'model_name': 'ensemble', 'outputs': [2, 2, 0, 2, 0], 'id': '0ebcc5f6f4c24d4d83eb36391eaefb98'}\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'model_name': 'ensemble',\n",
       " 'outputs': [2, 2, 0, 2, 0],\n",
       " 'id': '0ebcc5f6f4c24d4d83eb36391eaefb98'}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "server = function.to_mock_server()\n",
    "resp = server.test(\"/v2/models/infer\", body={\"inputs\": x})\n",
    "server.wait_for_completion()\n",
    "resp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Deploy the graph as a real-time serverless function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> 2021-01-09 22:49:40,088 [info] Starting remote function deploy\n",
      "2021-01-09 22:49:40  (info) Deploying function\n",
      "2021-01-09 22:49:40  (info) Building\n",
      "2021-01-09 22:49:40  (info) Staging files and preparing base images\n",
      "2021-01-09 22:49:40  (info) Building processor image\n",
      "2021-01-09 22:49:41  (info) Build complete\n",
      "2021-01-09 22:49:47  (info) Function deploy complete\n",
      "> 2021-01-09 22:49:48,422 [info] function deployed, address=default-tenant.app.yh55.iguazio-cd0.com:32222\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'http://default-tenant.app.yh55.iguazio-cd0.com:32222'"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function.deploy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<b>Invoke the remote function using the test data</b>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'model_name': 'ensemble',\n",
       " 'outputs': [1, 2, 0, 0, 0],\n",
       " 'id': '0ebcc5f6f4c24d4d83eb36391eaefb98'}"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function.invoke(\"/v2/models/infer\", body={\"inputs\": x})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
